`timescale 1ns/1ps
`include "../define.sv"

module pe_new_tb;

// 时钟和复位信号
reg clk;
reg reset;

// PE输入信号
reg [1:0] data_type;
reg mult_en;
reg c_add_en;
reg [31:0] a_in;
reg [31:0] b_in;
reg [31:0] c_in;
reg mix_precision;
reg pe_reset;

// PE输出信号
wire [31:0] a_out;
wire [31:0] b_out;
wire [31:0] c_out;
wire add_valid_o;
wire calc_done;
wire mult_en_out;

// 实例化被测模块
pe_new uut (
    .clk(clk),
    .reset(reset),
    .data_type(data_type),
    .mult_en(mult_en),
    .c_add_en(c_add_en),
    .a_in(a_in),
    .b_in(b_in),
    .c_in(c_in),
    .a_out(a_out),
    .b_out(b_out),
    .c_out(c_out),
    .mix_precision(mix_precision),
    .add_valid_o(add_valid_o),
    .calc_done(calc_done),
    .mult_en_out(mult_en_out),
    .pe_reset(pe_reset)
);

// 时钟生成
initial begin
    clk = 0;
    forever #5 clk = ~clk;
end

// 测试序列
initial begin
    // 初始化
    reset = 1;
    data_type = 2'b01;  // INT8
    mult_en = 0;
    c_add_en = 0;
    a_in = 32'h0;
    b_in = 32'h0;
    c_in = 32'h0;
    mix_precision = 0;
    pe_reset = 0;
    
    // 释放复位
    #20 reset = 0;
    
    // 测试1: 基本状态机转换
    $display("=== Test 1: Basic State Machine Transitions ===");
    
    // 等待几个时钟周期
    repeat(5) @(posedge clk);
    
    // 触发C矩阵加载，应该从IDLE转到CALCULATE
    $display("[%t] Triggering c_add_en", $time);
    c_add_en = 1;
    c_in = 32'h12345678;
    @(posedge clk);
    c_add_en = 0;
    
    // 等待计算完成
    repeat(10) @(posedge clk);
    
    // 模拟add_valid_o的下降沿
    $display("[%t] Simulating add_valid_o falling edge", $time);
    // 这里需要等待实际的add_valid_o信号
    
    // 等待DONE状态
    repeat(5) @(posedge clk);
    
    // 触发PE复位，应该从DONE转到IDLE
    $display("[%t] Triggering pe_reset", $time);
    pe_reset = 1;
    @(posedge clk);
    pe_reset = 0;
    
    // 等待复位完成
    repeat(5) @(posedge clk);
    
    // 测试2: 连续计算
    $display("=== Test 2: Continuous Calculation ===");
    
    // 再次触发计算
    c_add_en = 1;
    c_in = 32'h87654321;
    @(posedge clk);
    c_add_en = 0;
    
    // 等待计算完成
    repeat(15) @(posedge clk);
    
    // 再次复位
    pe_reset = 1;
    @(posedge clk);
    pe_reset = 0;
    
    // 等待复位完成
    repeat(5) @(posedge clk);
    
    $display("=== All Tests Completed ===");
    $finish;
end

// 监控状态变化
always @(posedge clk) begin
    $display("[%t] State: %s, c_add_en: %b, add_valid_o: %b, calc_done: %b, pe_reset: %b", 
             $time, uut.current_state.name(), c_add_en, add_valid_o, calc_done, pe_reset);
end

// 监控输出信号
always @(posedge clk) begin
    if (add_valid_o) begin
        $display("[%t] add_valid_o asserted, c_out: %h", $time, c_out);
    end
    
    if (calc_done) begin
        $display("[%t] calc_done asserted", $time);
    end
end

endmodule 